'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: with_support_function_call.py
* @Time: 2025/11/21
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain_core.messages import ToolMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

from adaptor.convert_to_openai_tool import convert_openai_tool
from mcp_core.mcp_client import MCPClient
from mcp_core.mcp_protocol import StatusCode

dotenv.load_dotenv()

with MCPClient(host="127.0.0.1", port=8080) as client:
    # 1、获取所有函数列表
    response = client.list_tools()
    if response.status_code == StatusCode.SUCCESS.value:
        tools = response.data["tools"]
    else:
        exit(0)
    openai_tools = convert_openai_tool(tools)
    #2、创建prompt
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是由OpenAI研发的聊天机器人, 可以帮助用户回答问题, 必要时候请调用工具帮助用户解答问题！"),
        ("human", "{query}")
    ])

    #3、创建大预言模型并绑定
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    llm_with_tool = llm.bind(tools=openai_tools)

    #4、创建链应用
    chain = {"query":RunnablePassthrough()} | prompt | llm_with_tool

    #5、解析输出
    query = ("请问现在是几点钟了？")

    resp = chain.invoke(query)
    tool_calls = resp.tool_calls # 拿到了  tool_calls 信息

    #6、判断是工具还是正常输出结果
    if len(tool_calls)<=0:#r如果 tool——calls 长度 不等于0 说明有函数调用信息
        print("生成了内容:", resp.content)# 没有 工具调用直接返回AI生成的内容
    else:# 处理调用信息
        #7、将 历史消息、人类消息、AI消息组合
        messages = prompt.invoke(query).to_messages()# 人类消息
        messages.append(resp)# AI 消息
        #8、循环遍历所有工具调用信息
        for tool_call in tool_calls: # tool_calls 是一个数组遍历调用所有 工具
            name = tool_call.get("name")#根据 AI 返回的 消息获取工具名称并提取工具实例
            print("正在执行工具:", name)
            id = tool_call.get("id")#tool message的id
            content = client.call_tool(name, tool_call.get("args")).data["result"]# 直接调用工具， 参数是从AI 返回的调用信息获取的
            print("工具输出:", content)
            messages.append(ToolMessage(# 合成工具消息， 让AI 总结
                content=content,
                tool_call_id = id,
            ))
            print("输出内容: ", llm.invoke(messages).content)
